# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: jhkim

from absl import app, flags

import collections
import datetime
import functools
import logging
import pandas
import numpy

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

from coin.base.book.types import TradeSide
from coin.strategy.mm import feed
from coin.strategy.mm.dumper_base import run_from_archive
from coin.exchange.okex_futures.kr_rest.futures_product import OkexFuturesProduct
from coin.exchange.kr_rest.product.product_impl import generate_product_from_str
from coin.strategy.mm.tool.kuona import KuonaInteractive, get_large_qty
from coin.exchange.util.feed_checker import get_default_feed_checker_config


class FeedAlignStrategy(object):
  def __init__(self, feed_product):
    self.feed_product = feed_product
    self.num_levels = 10
    self.kuona_cursor = KuonaInteractive(symbol=feed_product.symbol,
                                         large_order_qty=get_large_qty(feed_product))

  def on_book_reset(self, book_builder_name, book_builder):
    for product in [self.feed_product]:
      subscription_key = feed.get_subscription_key(product)
      if book_builder_name == subscription_key:
        book_builder.subscribe(product, functools.partial(self.on_book, product))
        book_builder.subscribe_trade(product, functools.partial(self.on_trade, product))

  def on_trade(self, product, trade):
    if product != self.feed_product:
      return
    if trade.side == TradeSide.TRADE_BUY_SIDE:
      trade_qty = trade.qty
    elif trade.side == TradeSide.TRADE_SELL_SIDE:
      trade_qty = -trade.qty
    else:
      raise ValueError(trade.side)
    self.kuona_cursor.trd_tpq.append((trade.timestamp, trade.price, trade_qty))

  def on_book(self, product, book):
    if product != self.feed_product:
      return
    if not hasattr(self, 'last_book_time'):
      self.last_book_time = 0
    if self.last_book_time + 2e8 > book.timestamp:
      return
    self.last_book_time = book.timestamp

    if product != self.feed_product:
      return
    if book.ask0().price <= book.bid0().price:
      return

    fields = collections.OrderedDict()
    n = self.num_levels

    bnames = ["bid%d" % (n - i - 1) for i in range(n)] + \
        ["ask%d" % i for i in range(n)]
    bnamemap = {}
    for i, bname in enumerate(bnames):
      bnamemap[bname] = "%s_%s" % (i, bname)

    timestr = pandas.DatetimeIndex([book.timestamp]).strftime('%Y%m%d %H:%M:%S')[0]
    fields['time'] = timestr

    orderbook = {}
    for side, method in [("bid", book.iteritems_bid), ("ask", book.iteritems_ask)]:
      for i, (price, price_qty) in enumerate(method()):
        if i >= n:
          break
        iname = "%s%d" % (side, i)
        name = bnamemap[iname]
        fields[iname + 'p'] = price
        fields[iname + 'q'] = price_qty.qty
        orderbook[price] = price_qty.qty

    self.kuona_cursor.rows.append(fields)
    self.kuona_cursor.obs.append(orderbook)
    self.kuona_cursor.tss.append(book.timestamp)
    self.kuona_cursor.ask0s.append(fields['ask0p'])
    self.kuona_cursor.bid0s.append(fields['bid0p'])

  def plot(self, ax):
    self.kuona_cursor.ax = ax
    self.kuona_cursor.cursor = len(self.kuona_cursor.tss) - flags.FLAGS.window
    self.kuona_cursor.window_size = flags.FLAGS.window
    self.kuona_cursor.plot_canvas()


def main(_):
  trading_date = flags.FLAGS.trading_date
  exchange = flags.FLAGS.exchange
  symbol = flags.FLAGS.symbol

  start_time = datetime.datetime.strptime(trading_date, "%Y%m%d")

  if flags.FLAGS.market_type == 'Futures':
    if flags.FLAGS.exchange == 'Okex':
      product = OkexFuturesProduct.FromStr(symbol, start_time)
    else:
      product = generate_product_from_str(flags.FLAGS.exchange, None, symbol)
  else:
    raise ValueError(flags.FLAGS.market_type)

  hours = 1

  fig = plt.figure(figsize=(18, 9))

  prev_ax = None

  min_tick = None
  max_tick = None

  axs = []

  for i, worker_id in enumerate([['1', '2'], '1', '2']):
    use_feed_arbitration = type(worker_id) is list
    ax = fig.add_subplot(3, 1, i + 1, sharex=prev_ax)
    strategy = FeedAlignStrategy(product)
    run_from_archive([strategy.feed_product],
                     strategy.on_book_reset,
                     start_time,
                     start_time + datetime.timedelta(hours=hours),
                     machine=[flags.FLAGS.feed_machine],
                     root_dir=flags.FLAGS.feed_root,
                     use_feed_arbitration=use_feed_arbitration,
                     worker_id=worker_id,
                     feed_checker_config=get_default_feed_checker_config([strategy.feed_product]))
    strategy.plot(ax)

    axs.append(ax)

    ax.set_title(str(worker_id))
    print(worker_id, strategy.kuona_cursor.tss[-flags.FLAGS.window:])
    if prev_ax is None:
      prev_ax = ax

  ticks = ax.get_xticks()
  min_tick = min_tick or ticks[0]
  min_tick = min(min_tick, ticks[0])
  max_tick = max_tick or ticks[-1]
  max_tick = max(max_tick, ticks[-1])
  max_tick += (ticks[-1] - ticks[-2])

  for ax in axs:
    minor_ticks = numpy.linspace(min_tick, max_tick, 200)
    ax.set_xticks(minor_ticks, minor=True)
    ax.grid(color='k', linestyle='--', linewidth=0.5, which='minor')

  plt.savefig("compare_feed.png")
  plt.close()


if __name__ == '__main__':
  flags.DEFINE_string('trading_date', '20180820', 'Trading date in form of %Y%m%d.')

  flags.DEFINE_string('feed_machine', 'feed-01.ap-northeast-1.aws', 'feed machine name.')

  flags.DEFINE_string('feed_root', None, 'feed root')

  flags.DEFINE_string('exchange', 'Okex', 'Exchange name.')

  flags.DEFINE_string('market_type', 'Futures', 'Spot, Futures')

  flags.DEFINE_string('symbol', 'BTC-USD.QUARTER', 'symbol name')

  flags.DEFINE_integer('window', 10, '')

  logging.basicConfig(level='DEBUG', format='%(levelname)8s %(asctime)s %(name)s] %(message)s')
  app.run(main)
